package weiyao.GUI;

import javax.swing.*;
import javax.swing.border.TitledBorder;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.List;
import us.codecraft.webmagic.Spider;
import weiyao.CrawlerProcessor.CrawlerProcessor;
import weiyao.DatabaseUtil;
import weiyao.SiteManager;
import weiyao.Vul.VulnerabilityConfig;
import weiyao.Vul.VulnerabilityScanner;

public class ScannerGUI extends JFrame implements ScanResultCallback {
    // GUI组件
    private JTextField urlField;
    private JButton startButton, stopButton, clearLogButton, saveToDBButton, generateReportButton;
    private JTextArea logArea;
    private JProgressBar progressBar;
    private JTabbedPane resultTabs;
    private DefaultListModel<String> vulnerabilityListModel;
    private JList<String> vulnerabilityList;
    private JLabel sqlLabel, xssLabel, csrfLabel, wpLabel;
    private DefaultListModel<String> pageListModel;

    // 漏洞复选框
    private final Map<VulnerabilityConfig.VulnerabilityType, JCheckBox> vulCheckboxes =
            new EnumMap<>(VulnerabilityConfig.VulnerabilityType.class);

    // 扫描状态
    private boolean isScanning = false;
    private Spider spider;

    // 漏洞数据
    private int sqlInjectionCount = 0;
    private int xssCount = 0;
    private int csrfCount = 0;
    private int weakPasswordCount = 0;
    private List<String> scannedUrls = new ArrayList<>();
    private List<String> vulnerabilityDetails = new ArrayList<>();
    private final List<VulnerabilityRecord> vulnerabilityRecords = new ArrayList<>();

    // 管理器
    private final SiteManager siteManager = new SiteManager();
    private VulnerabilityScanner vulnerabilityScanner;

    public ScannerGUI() {
        initUI();
        setTitle("Web漏洞扫描器");
        setSize(900, 700);
        setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setLocationRelativeTo(null);

        // 添加事件监听器
        startButton.addActionListener(this::startScan);
        stopButton.addActionListener(e -> stopScan());
        clearLogButton.addActionListener(e -> logArea.setText(""));
        saveToDBButton.addActionListener(this::saveToDatabase);
        generateReportButton.addActionListener(this::generateReport);
    }

    private void initUI() {
        // 主面板
        JPanel mainPanel = new JPanel(new BorderLayout(10, 10));
        mainPanel.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10));

        // 顶部输入面板
        JPanel inputPanel = new JPanel(new BorderLayout(5, 5));
        urlField = new JTextField("http://192.168.190.139:85/", 30);
        startButton = new JButton("开始扫描");
        stopButton = new JButton("停止扫描");
        clearLogButton = new JButton("清空日志");
        saveToDBButton = new JButton("保存到数据库");
        generateReportButton = new JButton("生成报告");

        stopButton.setEnabled(false);
        saveToDBButton.setEnabled(false);
        generateReportButton.setEnabled(false);

        // 扫描选项面板
        JPanel scanOptionsPanel = new JPanel(new FlowLayout(FlowLayout.LEFT));
        scanOptionsPanel.setBorder(BorderFactory.createTitledBorder("扫描选项"));
        scanOptionsPanel.add(new JLabel("检测类型:"));

        // 添加漏洞复选框
        for (VulnerabilityConfig.VulnerabilityType type :
                VulnerabilityConfig.getSupportedTypes()) {

            String label = "";
            switch (type) {
                case SQL_INJECTION: label = "SQL注入"; break;
                case XSS: label = "XSS"; break;
                case CSRF: label = "CSRF"; break;
                case WEAK_PASSWORD: label = "弱口令"; break;
            }

            JCheckBox checkbox = new JCheckBox(label, true);
            scanOptionsPanel.add(checkbox);
            vulCheckboxes.put(type, checkbox);
        }

        JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT, 5, 0));
        buttonPanel.add(startButton);
        buttonPanel.add(stopButton);
        buttonPanel.add(clearLogButton);
        buttonPanel.add(saveToDBButton);
        buttonPanel.add(generateReportButton);

        inputPanel.add(new JLabel("目标URL:"), BorderLayout.WEST);
        inputPanel.add(urlField, BorderLayout.CENTER);
        inputPanel.add(buttonPanel, BorderLayout.EAST);

        // 添加扫描选项面板
        JPanel topPanel = new JPanel(new BorderLayout());
        topPanel.add(inputPanel, BorderLayout.NORTH);
        topPanel.add(scanOptionsPanel, BorderLayout.SOUTH);

        // 日志区域
        logArea = new JTextArea();
        logArea.setEditable(false);
        logArea.setFont(new Font("Monospaced", Font.PLAIN, 12));
        JScrollPane logScroll = new JScrollPane(logArea);
        logScroll.setBorder(BorderFactory.createTitledBorder("扫描日志"));

        // 进度条
        progressBar = new JProgressBar();
        progressBar.setStringPainted(true);
        progressBar.setBorder(BorderFactory.createEmptyBorder(5, 0, 0, 0));

        // 结果面板
        resultTabs = new JTabbedPane();
        initResultTabs();

        // 添加组件到主面板
        mainPanel.add(topPanel, BorderLayout.NORTH);
        mainPanel.add(logScroll, BorderLayout.CENTER);
        mainPanel.add(progressBar, BorderLayout.SOUTH);

        // 添加主面板和结果面板
        JSplitPane splitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT, mainPanel, resultTabs);
        splitPane.setResizeWeight(0.6);
        splitPane.setDividerLocation(400);
        add(splitPane);
    }

    private void initResultTabs() {
        JPanel resultPanel = new JPanel(new BorderLayout());

        // 漏洞列表
        vulnerabilityListModel = new DefaultListModel<>();
        vulnerabilityList = new JList<>(vulnerabilityListModel);
        vulnerabilityList.setSelectionMode(ListSelectionModel.SINGLE_SELECTION);
        JScrollPane listScroll = new JScrollPane(vulnerabilityList);
        listScroll.setBorder(BorderFactory.createTitledBorder("检测到的漏洞"));

        // 漏洞详情
        JTextArea detailArea = new JTextArea();
        detailArea.setEditable(false);
        detailArea.setFont(new Font("Monospaced", Font.PLAIN, 12));
        JScrollPane detailScroll = new JScrollPane(detailArea);
        detailScroll.setBorder(BorderFactory.createTitledBorder("漏洞详情"));
        detailArea.setLineWrap(true);
        detailArea.setWrapStyleWord(true);
        vulnerabilityList.addListSelectionListener(e -> {
            if (!e.getValueIsAdjusting()) {
                int index = vulnerabilityList.getSelectedIndex();
                if (index >= 0) {
                    detailArea.setText(vulnerabilityDetails.get(index));
                }
            }
        });

        // 统计面板
        JPanel statsPanel = new JPanel(new GridLayout(4, 1, 5, 5));
        statsPanel.setBorder(BorderFactory.createTitledBorder("漏洞统计"));

        sqlLabel = new JLabel("SQL注入: 0");
        xssLabel = new JLabel("XSS漏洞: 0");
        csrfLabel = new JLabel("CSRF漏洞: 0");
        wpLabel = new JLabel("弱口令: 0");

        statsPanel.add(sqlLabel);
        statsPanel.add(xssLabel);
        statsPanel.add(csrfLabel);
        statsPanel.add(wpLabel);

        // 布局
        JSplitPane splitPane = new JSplitPane(JSplitPane.HORIZONTAL_SPLIT, listScroll, detailScroll);
        splitPane.setDividerLocation(300);

        resultPanel.add(splitPane, BorderLayout.CENTER);
        resultPanel.add(statsPanel, BorderLayout.EAST);

        resultTabs.addTab("扫描结果", resultPanel);

        // 扫描页面标签页
        pageListModel = new DefaultListModel<>();
        JList<String> pageList = new JList<>(pageListModel);
        JScrollPane pageScroll = new JScrollPane(pageList);
        pageScroll.setBorder(BorderFactory.createTitledBorder("已扫描页面"));
        resultTabs.addTab("扫描页面", pageScroll);
    }

    private void startScan(ActionEvent e) {
        if (isScanning) return;

        String url = urlField.getText().trim();
        if (url.isEmpty()) {
            JOptionPane.showMessageDialog(this, "请输入有效的URL", "错误", JOptionPane.ERROR_MESSAGE);
            return;
        }

        // 获取启用的扫描类型
        Set<VulnerabilityConfig.VulnerabilityType> enabledScans = new HashSet<>();
        for (Map.Entry<VulnerabilityConfig.VulnerabilityType, JCheckBox> entry :
                vulCheckboxes.entrySet()) {

            if (entry.getValue().isSelected()) {
                enabledScans.add(entry.getKey());
            }
        }

        if (enabledScans.isEmpty()) {
            JOptionPane.showMessageDialog(this, "请至少选择一种扫描类型", "错误", JOptionPane.ERROR_MESSAGE);
            return;
        }

        // 配置站点
        String domain = extractDomain(url);
        siteManager.configureSite(domain);

        // 创建扫描器
        vulnerabilityScanner = new VulnerabilityScanner(enabledScans, siteManager);

        log("开始扫描: " + url);
        isScanning = true;
        startButton.setEnabled(false);
        stopButton.setEnabled(true);
        saveToDBButton.setEnabled(false);
        generateReportButton.setEnabled(false);
        progressBar.setValue(0);
        progressBar.setString("扫描中...");
        vulnerabilityListModel.clear();
        vulnerabilityDetails.clear();
        pageListModel.clear();
        vulnerabilityRecords.clear();
        sqlInjectionCount = 0;
        xssCount = 0;
        csrfCount = 0;
        weakPasswordCount = 0;
        scannedUrls.clear();
        updateStats();

        // 创建爬虫处理器
        CrawlerProcessor processor = new CrawlerProcessor(
                siteManager,
                vulnerabilityScanner,
                this,
                url
        );

        // 启动爬虫
        spider = Spider.create(processor)
                .addUrl(url)
                .thread(5);

        new Thread(() -> {
            spider.run();
            SwingUtilities.invokeLater(() -> {
                stopScan();
                onScanComplete();
                saveToDBButton.setEnabled(true);
                generateReportButton.setEnabled(true);
            });
        }).start();
    }

    private String extractDomain(String url) {
        try {
            url = url.replaceFirst("^(https?://)", "");
            int endIndex = url.indexOf('/');
            if (endIndex != -1) {
                return url.substring(0, endIndex);
            }
            return url;
        } catch (Exception e) {
            return "unknown-domain";
        }
    }

    private void stopScan() {
        if (spider != null) {
            spider.stop();
        }
        isScanning = false;
        startButton.setEnabled(true);
        stopButton.setEnabled(false);
        progressBar.setString("已停止");
        log("扫描已停止");
    }

    @Override
    public void onVulnerabilityFound(String type, String details) {
        SwingUtilities.invokeLater(() -> {
            // 更新计数
            switch (type) {
                case "SQL注入": sqlInjectionCount++; break;
                case "XSS": xssCount++; break;
                case "CSRF": csrfCount++; break;
                case "弱口令": weakPasswordCount++; break;
            }
            updateStats();

            // 打印日志
            log("发现" + type + "漏洞: " + details);

            // 创建简短的漏洞摘要显示在列表中
            String summary = createVulnerabilitySummary(type, details);

            // 添加到漏洞列表
            vulnerabilityListModel.addElement(summary);
            vulnerabilityDetails.add(details);

            // 创建并存储漏洞记录
            VulnerabilityRecord record = parseVulnerabilityDetails(type, details);
            if (record != null) {
                vulnerabilityRecords.add(record);
            }
        });
    }

    // 创建漏洞摘要（在列表中显示）
    private String createVulnerabilitySummary(String type, String details) {
        StringBuilder summary = new StringBuilder(type);

        // 对于SQL注入，添加提取的登录信息摘要
        if ("SQL注入".equals(type)) {
            String loginInfo = extractInfo(details, "提取的登录信息");
            String passwordInfo = extractInfo(details, "提取的密码信息");

            if (!loginInfo.isEmpty()) {
                summary.append(" - 登录名: ").append(truncate(loginInfo, 20));
            }
            if (!passwordInfo.isEmpty()) {
                summary.append(" - 密码: ").append(truncate(passwordInfo, 20));
            }
        }

        return summary.toString();
    }

    // 提取特定信息（辅助方法）
    private String extractInfo(String details, String key) {
        if (details.contains(key + ": ")) {
            int start = details.indexOf(key + ": ") + key.length() + 2;
            int end = details.indexOf("\n", start);
            if (end == -1) end = details.length();
            return details.substring(start, end).trim();
        }
        return "";
    }

    // 截断长文本（辅助方法）
    private String truncate(String text, int maxLength) {
        if (text.length() <= maxLength) return text;
        return text.substring(0, maxLength) + "...";
    }

    // 解析漏洞详情
    private VulnerabilityRecord parseVulnerabilityDetails(String type, String details) {
        VulnerabilityRecord record = new VulnerabilityRecord();
        record.setType(type);

        // 尝试从详情中提取URL
        String url = extractUrlFromDetails(details);
        if (url != null) {
            record.setUrl(url);
        }

        // 解析其他关键信息
        Map<String, String> infoMap = new HashMap<>();

        // 根据漏洞类型使用不同的解析逻辑
        switch (type) {
            case "SQL注入":
                extractInfo(details, "参数", infoMap);
                extractInfo(details, "载荷", infoMap);
                // 添加提取的登录名和密码信息
                extractInfo(details, "提取的登录信息", infoMap);
                extractInfo(details, "提取的密码信息", infoMap);
                break;

            case "XSS":
                extractInfo(details, "参数", infoMap);
                extractInfo(details, "载荷", infoMap);
                extractInfo(details, "上下文", infoMap);
                break;

            case "CSRF":
                extractInfo(details, "URL", infoMap);
                extractInfo(details, "方法", infoMap);
                extractInfo(details, "载荷", infoMap);
                break;

            case "弱口令":
                extractInfo(details, "用户名", infoMap);
                extractInfo(details, "密码", infoMap);
                break;
        }

        record.setDetails(infoMap);
        return record;
    }

    private String extractUrlFromDetails(String details) {
        // 尝试从常见模式中提取URL
        if (details.contains("URL: ")) {
            int start = details.indexOf("URL: ") + 5;
            int end = details.indexOf("\n", start);
            if (end == -1) end = details.length();
            return details.substring(start, end).trim();
        }
        return null;
    }

    private void extractInfo(String details, String key, Map<String, String> infoMap) {
        if (details.contains(key + ": ")) {
            int start = details.indexOf(key + ": ") + key.length() + 2;
            int end = details.indexOf("\n", start);
            if (end == -1) end = details.length();
            infoMap.put(key, details.substring(start, end).trim());
        }
    }

    private void updateStats() {
        sqlLabel.setText("SQL注入: " + sqlInjectionCount);
        xssLabel.setText("XSS漏洞: " + xssCount);
        csrfLabel.setText("CSRF漏洞: " + csrfCount);
        wpLabel.setText("弱口令: " + weakPasswordCount);
    }

    @Override
    public void onPageScanned(String url) {
        SwingUtilities.invokeLater(() -> {
            scannedUrls.add(url);
            log("扫描页面: " + url);
            pageListModel.addElement(url);

            // 更新进度
            progressBar.setString("已扫描: " + scannedUrls.size() + " 页面");
        });
    }

    @Override
    public void onError(String message) {
        SwingUtilities.invokeLater(() -> {
            log("错误: " + message);
        });
    }

    @Override
    public void onStatusUpdate(String message) {
        SwingUtilities.invokeLater(() -> log(message));
    }

    @Override
    public void onScanComplete() {
        SwingUtilities.invokeLater(() -> {
            log("扫描完成!");
            log("扫描页面数: " + scannedUrls.size());
            log("发现漏洞: " + vulnerabilityListModel.size());

            progressBar.setValue(100);
            progressBar.setString("扫描完成");

            JOptionPane.showMessageDialog(this,
                    "扫描完成!\n" +
                            "扫描页面: " + scannedUrls.size() + "\n" +
                            "发现漏洞: " + vulnerabilityListModel.size() + "\n" +
                            "- SQL注入: " + sqlInjectionCount + "\n" +
                            "- XSS: " + xssCount + "\n" +
                            "- CSRF: " + csrfCount + "\n" +
                            "- 弱口令: " + weakPasswordCount,
                    "扫描结果",
                    JOptionPane.INFORMATION_MESSAGE);
        });
    }

    private void log(String message) {
        SwingUtilities.invokeLater(() -> {
            logArea.append("[" + new java.util.Date() + "] " + message + "\n");
            logArea.setCaretPosition(logArea.getDocument().getLength());
        });
    }

    private void saveToDatabase(ActionEvent e) {
        if (vulnerabilityRecords.isEmpty()) {
            JOptionPane.showMessageDialog(this, "没有漏洞记录可保存", "提示", JOptionPane.INFORMATION_MESSAGE);
            return;
        }

        new SwingWorker<Void, Void>() {
            @Override
            protected Void doInBackground() throws Exception {
                DatabaseUtil.saveVulnerabilities(vulnerabilityRecords);
                return null;
            }

            @Override
            protected void done() {
                try {
                    get(); // 检查是否有异常
                    JOptionPane.showMessageDialog(ScannerGUI.this,
                            "成功保存 " + vulnerabilityRecords.size() + " 条漏洞记录",
                            "保存成功", JOptionPane.INFORMATION_MESSAGE);
                } catch (Exception ex) {
                    JOptionPane.showMessageDialog(ScannerGUI.this,
                            "保存失败: " + ex.getMessage(),
                            "数据库错误", JOptionPane.ERROR_MESSAGE);
                }
            }
        }.execute();
    }

    // 生成报告功能
    private void generateReport(ActionEvent e) {
        if (vulnerabilityRecords.isEmpty()) {
            JOptionPane.showMessageDialog(this, "没有漏洞记录可生成报告", "提示", JOptionPane.INFORMATION_MESSAGE);
            return;
        }

        JFileChooser fileChooser = new JFileChooser();
        fileChooser.setDialogTitle("保存报告");
        fileChooser.setSelectedFile(new File("漏洞扫描报告_" + new SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()) + ".txt"));

        int userSelection = fileChooser.showSaveDialog(this);
        if (userSelection != JFileChooser.APPROVE_OPTION) {
            return;
        }

        File reportFile = fileChooser.getSelectedFile();

        new SwingWorker<Void, Void>() {
            @Override
            protected Void doInBackground() throws Exception {
                try (BufferedWriter writer = new BufferedWriter(new FileWriter(reportFile))) {
                    // 生成报告头部
                    writer.write("========================================");
                    writer.newLine();
                    writer.write("         Web漏洞扫描报告");
                    writer.newLine();
                    writer.write("========================================");
                    writer.newLine();
                    writer.newLine();

                    // 扫描信息
                    writer.write("扫描信息:");
                    writer.newLine();
                    writer.write("----------------------------------------");
                    writer.newLine();
                    writer.write("扫描时间: " + new Date());
                    writer.newLine();
                    writer.write("目标URL: " + urlField.getText());
                    writer.newLine();
                    writer.write("扫描页面数: " + scannedUrls.size());
                    writer.newLine();
                    writer.write("发现漏洞总数: " + vulnerabilityRecords.size());
                    writer.newLine();
                    writer.newLine();

                    // 漏洞统计
                    writer.write("漏洞统计:");
                    writer.newLine();
                    writer.write("----------------------------------------");
                    writer.newLine();
                    writer.write("SQL注入: " + sqlInjectionCount);
                    writer.newLine();
                    writer.write("XSS漏洞: " + xssCount);
                    writer.newLine();
                    writer.write("CSRF漏洞: " + csrfCount);
                    writer.newLine();
                    writer.write("弱口令: " + weakPasswordCount);
                    writer.newLine();
                    writer.newLine();

                    // 漏洞详情
                    writer.write("漏洞详情:");
                    writer.newLine();
                    writer.write("----------------------------------------");
                    writer.newLine();

                    for (int i = 0; i < vulnerabilityRecords.size(); i++) {
                        VulnerabilityRecord record = vulnerabilityRecords.get(i);
                        writer.write((i + 1) + ". [" + record.getType() + "]");
                        writer.newLine();

                        if (record.getUrl() != null) {
                            writer.write("URL: " + record.getUrl());
                            writer.newLine();
                        }

                        for (Map.Entry<String, String> entry : record.getDetails().entrySet()) {
                            writer.write(entry.getKey() + ": " + entry.getValue());
                            writer.newLine();
                        }

                        writer.newLine();
                    }

                    // 扫描页面列表
                    writer.write("扫描页面列表:");
                    writer.newLine();
                    writer.write("----------------------------------------");
                    writer.newLine();
                    for (String url : scannedUrls) {
                        writer.write(url);
                        writer.newLine();
                    }

                    writer.newLine();
                    writer.write("========================================");
                    writer.newLine();
                    writer.write("报告生成时间: " + new Date());
                    writer.newLine();
                }
                return null;
            }

            @Override
            protected void done() {
                try {
                    get(); // 检查是否有异常
                    JOptionPane.showMessageDialog(ScannerGUI.this,
                            "报告已成功生成: " + reportFile.getAbsolutePath(),
                            "报告生成成功", JOptionPane.INFORMATION_MESSAGE);
                } catch (Exception ex) {
                    JOptionPane.showMessageDialog(ScannerGUI.this,
                            "生成报告失败: " + ex.getMessage(),
                            "错误", JOptionPane.ERROR_MESSAGE);
                }
            }
        }.execute();
    }
}